import gradio as gr
import csv
import joblib
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
import jieba
import warnings
warnings.filterwarnings('ignore')
import numpy as np
import torch
from transformers import BertTokenizer, BertModel
from Classifier import Classifier
import warnings
warnings.filterwarnings("ignore")

model_name = '../bert-base-chinese'
model_path = '../bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
bert_model = BertModel.from_pretrained(model_path)
model = Classifier(bert_model)
# 加载最佳模型的权重
路径='saved_weights_外卖.pt'
model.load_state_dict(torch.load(路径))
def image_to_text(text):
    sent_id = tokenizer.encode(text,
                               add_special_tokens=True,
                               max_length=100,
                               truncation=True,
                               pad_to_max_length='right')

    att_mask = [int(tok > 0) for tok in sent_id]
    sent_id = torch.tensor(sent_id)
    att_mask = torch.tensor(att_mask)
    sent_id = sent_id.unsqueeze(0)
    att_mask = att_mask.unsqueeze(0)
    preds = model(sent_id, att_mask)

    total_preds = []
    total_preds.append(np.argmax(preds.detach().cpu().numpy(), axis=1))
    if 1 in np.concatenate(total_preds, axis=0):
        return "好评"
    else:
        return "差评"




demo = gr.Interface(fn=image_to_text, inputs="text", outputs="text", title="评论舆情判别",
                    description="评论舆情判别")
demo.launch(server_name='127.0.0.1', server_port=6001)
